import ptan
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical

class MPOAgent(ptan.agent.BaseAgent):
    def __init__(self, model, action_selector=ptan.actions.ProbabilityActionSelector(), device="cpu",
                 apply_softmax=False, preprocessor=ptan.agent.default_states_preprocessor):
        self.model = model
        self.action_selector = action_selector
        self.device = device
        self.apply_softmax = apply_softmax
        self.preprocessor = preprocessor

    @torch.no_grad()
    def __call__(self, states, agent_states=None):

        if self.preprocessor is not None:
            states = self.preprocessor(states)
            if torch.is_tensor(states):
                states = states.to(self.device)
        actions, probs = self.model.action(states)
        return np.array(actions.cpu().numpy()), [(probs)]



class FeatureExtractor(nn.Module):
    def __init__(self, input_shape):
        super(FeatureExtractor, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(input_shape[0], 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU()
        )


    def get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x):
        x = x.float() / 255.0
        return self.conv(x)
    

class LinearFeatureExtractor(nn.Module):
    def __init__(self, input_shape):
        super(LinearFeatureExtractor, self).__init__()

        self.conv = nn.Sequential(
            nn.Linear(input_shape, 256),
            nn.ReLU(),
        )

        self.out_size = 256


    def forward(self, x):
        return self.conv(x)
    

class DQNModel(nn.Module):
    def __init__(self, input_shape, n_actions):
        super(DQNModel, self).__init__()

        self.feature_extractor = LinearFeatureExtractor(input_shape)

        self.fc = nn.Sequential(
            nn.Linear(self.feature_extractor.out_size, 512),
            nn.ReLU(),
            nn.Linear(512, n_actions)
        )

    def forward(self, x):
        return self.fc(self.feature_extractor(x))
    

class AtariA2C(nn.Module):
    def __init__(self, input_shape, n_actions):
        super(AtariA2C, self).__init__()

        # obs_action = (input_shape[2], input_shape[0], input_shape[1])
        print("obs_action: ", input_shape)
        obs_action = input_shape

        self.feature_extractor = LinearFeatureExtractor(obs_action)

        self.policy = nn.Sequential(
            nn.Linear(self.feature_extractor.out_size, 512),
            nn.ReLU(),
            nn.Linear(512, n_actions)
        )

        self.value = nn.Sequential(
            nn.Linear(self.feature_extractor.out_size, 512),
            nn.ReLU(),
            nn.Linear(512, 1)
        )

    def forward(self, x):
        conv_out = self.feature_extractor(x)
        return self.policy(conv_out), self.value(conv_out)